﻿#include <cstdio>

extern "C" {
#include <libavutil/opt.h>
#include <libavutil/samplefmt.h>
#include <libavutil/channel_layout.h>
#include <libavutil/common.h>
#include <libavcodec/avcodec.h>
#include <libavformat/avformat.h>
#include <libavfilter/avfilter.h>
#include <libavfilter/buffersink.h>
#include <libavfilter/buffersrc.h>
}

#include <iostream>

int main()
{
    int ret = -1;
    const char* input_filename = "input.mp3";
    const char* output_filename = "output.png";

    // 初始化FFmpeg
    av_register_all();
    avcodec_register_all();
    avformat_network_init();

    AVFormatContext* format_ctx = avformat_alloc_context();
    if (avformat_open_input(&format_ctx, input_filename, nullptr, nullptr) != 0)
    {
        fprintf(stderr, "Error opening input file\n");
        return -1;
    }


    if (avformat_find_stream_info(format_ctx, nullptr) < 0)
    {
        fprintf(stderr, "Could not find stream information\n");
        return -1;
    }

    AVCodec* codec = avcodec_find_decoder(format_ctx->streams[0]->codecpar->codec_id);
    AVCodecContext* codec_ctx = avcodec_alloc_context3(codec);
    avcodec_parameters_to_context(codec_ctx, format_ctx->streams[0]->codecpar);
    avcodec_open2(codec_ctx, codec, nullptr);


    AVFilterContext* buffersrc_ctx = nullptr;
    AVFilterContext* showfreqs_ctx = nullptr;
    AVFilterContext* buffersink_ctx = nullptr;

    AVFilterGraph* pFilterGraph = avfilter_graph_alloc();

    const AVFilter* abuffer = avfilter_get_by_name("abuffer");
    const AVFilter* showfreqs = avfilter_get_by_name("showspectrumpic");
    const AVFilter* buffersink = avfilter_get_by_name("buffersink");

    char args[512];
    snprintf(args, sizeof(args),
        "time_base=%d/%d:sample_rate=%d:sample_fmt=%s:channel_layout=0x%" PRIx64,
        codec_ctx->time_base.num, codec_ctx->time_base.den, codec_ctx->sample_rate,
        av_get_sample_fmt_name(codec_ctx->sample_fmt), codec_ctx->channel_layout);

    const char* show_filter_desc = "s=640x480:scale=5thrt:legend=0:start=20:stop=16000";

    avfilter_graph_create_filter(&buffersrc_ctx, abuffer, "in", args, nullptr, pFilterGraph);
    avfilter_graph_create_filter(&showfreqs_ctx, showfreqs, "showspectrumpic", show_filter_desc, nullptr, pFilterGraph);
    avfilter_graph_create_filter(&buffersink_ctx, buffersink, "out", nullptr, nullptr, pFilterGraph);

    avfilter_link(buffersrc_ctx, 0, showfreqs_ctx, 0);
    avfilter_link(showfreqs_ctx, 0, buffersink_ctx, 0);

    avfilter_graph_config(pFilterGraph, nullptr);

    // 打开输出文件
    FILE* output_file = fopen(output_filename, "wb");

    AVPacket packet;
    av_init_packet(&packet);

    while (av_read_frame(format_ctx, &packet) >= 0)
    {
        if (packet.stream_index == 0)
        {
            if (avcodec_send_packet(codec_ctx, &packet) == 0)
            {
                AVFrame* outFrame = av_frame_alloc();
                while (avcodec_receive_frame(codec_ctx, outFrame) == 0)
                {
                    int numBytes = av_get_bytes_per_sample(codec_ctx->sample_fmt);

                    // 发送音频帧到滤镜图
                    av_buffersrc_add_frame(buffersrc_ctx, outFrame);
                    
                    // 获取输出帧
                    AVFrame* filteredFrame = av_frame_alloc();
                    ret = av_buffersink_get_frame(buffersink_ctx, filteredFrame);

                    if (ret == AVERROR(EAGAIN) || ret == AVERROR_EOF)
                    {
                        fprintf(stderr, "Error getting frame from filter, ret: %d\n", ret);
                        av_frame_free(&filteredFrame);
                        continue;
                    }
                    else if (ret < 0)
                    {
                        fprintf(stderr, "Error getting frame from filter, ret: %d\n", ret);
                        break;
                    }
                    else
                    {
                        // 将输出帧写入文件
                        fwrite(filteredFrame->data[0], 1, filteredFrame->linesize[0], output_file);

                    }

                    av_frame_free(&filteredFrame);
                    
                }

                av_frame_free(&outFrame);
            }
        }
        av_packet_unref(&packet);
    }


    fclose(output_file);
    avformat_close_input(&format_ctx);
    avcodec_free_context(&codec_ctx);
    avfilter_free(buffersrc_ctx);
    avfilter_free(buffersink_ctx);
    avfilter_graph_free(&pFilterGraph);

    avformat_network_deinit();
    avcodec_register_all();

    return 0;
}
